import numpy as np
import torch

class PrivacyBudgetManager:
    def __init__(self, total_budget, sensitivity_levels, delta=1e-5, min_budget_factor=0.1):
        """
        隐私预算管理模块 (对应论文2.5节)
        
        参数:
            total_budget: 总隐私预算 (ε_total)
            sensitivity_levels: 客户端敏感度等级列表
            delta: 差分隐私的δ参数
            min_budget_factor: 最低隐私预算比例
        """
        self.total_budget = total_budget
        self.sensitivity_levels = sensitivity_levels
        self.delta = delta
        self.min_budget_factor = min_budget_factor
        self.privacy_consumption = np.zeros_like(sensitivity_levels)
        self.round_consumption_history = []
        
    def dynamic_budget_allocation(self, gradient_norms, round_num):
        """
        隐私预算动态分配 (对应论文2.5.2节)
        
        参数:
            gradient_norms: 客户端梯度范数列表
            round_num: 当前轮次
            
        返回:
            budget_allocation: 各客户端分配的隐私预算
        """
        num_clients = len(self.sensitivity_levels)
        
        # 客户端敏感度分级 (敏感度越高，分配的隐私预算越低)
        sensitivity_weights = 1.0 / (self.sensitivity_levels + 1e-8)
        
        # 动态分配隐私预算 (基于敏感度和梯度范数)
        budget_allocation = np.zeros(num_clients)
        for i in range(num_clients):
            # 公式: budget ∝ (1/sensitivity) * gradient_norm
            budget_allocation[i] = sensitivity_weights[i] * gradient_norms[i]
        
        # 归一化并应用总预算约束
        budget_allocation = budget_allocation / budget_allocation.sum() * self.total_budget
        
        # 个体约束：确保最低隐私强度
        min_budget = self.min_budget_factor * self.total_budget / num_clients
        budget_allocation = np.maximum(budget_allocation, min_budget)
        
        return budget_allocation
    
    def calculate_noise_scale(self, clip_threshold, budget_allocation):
        """
        计算噪声尺度 (对应论文公式13)
        
        参数:
            clip_threshold: 梯度裁剪阈值 (C)
            budget_allocation: 分配的隐私预算
            
        返回:
            noise_scales: 各客户端的噪声尺度
        """
        # 高斯机制下的噪声标准差公式: σ = C * sqrt(2*log(1.25/δ)) / ε
        noise_scales = clip_threshold * np.sqrt(2 * np.log(1.25 / self.delta)) / (budget_allocation + 1e-8)
        return noise_scales
    
    def update_consumption(self, consumption):
        """
        更新隐私消耗
        参数:
            consumption: 本轮各客户端的隐私消耗
        """
        self.privacy_consumption += consumption
        self.round_consumption_history.append(consumption)
    
    def check_budget_constraint(self):
        """检查预算约束 (对应论文公式12)"""
        return np.all(self.privacy_consumption <= self.total_budget)
    
    def get_remaining_budget(self):
        """获取剩余隐私预算"""
        return self.total_budget - np.sum(self.privacy_consumption)
    
    def get_total_consumption(self):
        """获取总隐私消耗"""
        return np.sum(self.privacy_consumption)
